# 导入FastAPI和WebSockets
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse

# 创建一个FastAPI实例
app = FastAPI()

# 创建一个全局的字典，用于存储所有的WebSocket连接
connections = {}

# 创建一个HTML文档作为客户端的示例，可以在浏览器中运行
html = """
<!DOCTYPE html>
<html>
    <head>
        <title>Chat Room</title>
    </head>
    <body>
        <h1>Chat Room</h1>
        <form id="form">
            <input type="text" id="username" placeholder="Enter your username">
            <button id="connect" type="submit">Connect</button>
        </form>
        <div id="messages"></div>
        <form id="send_form" style="display: none">
            <input type="text" id="message" placeholder="Enter your message">
            <button id="send" type="submit">Send</button>
        </form>
        <script>
            // 获取HTML元素
            var form = document.getElementById("form")
            var connect = document.getElementById("connect")
            var username = document.getElementById("username")
            var messages = document.getElementById("messages")
            var send_form = document.getElementById("send_form")
            var send = document.getElementById("send")
            var message = document.getElementById("message")

            // 创建一个WebSocket对象，连接到服务器端的/ws路由
            var ws = new WebSocket("ws://localhost:8000/ws")

            // 当WebSocket连接建立时，执行以下操作
            ws.onopen = function(event) {
                // 隐藏连接表单，显示发送表单
                form.style.display = "none"
                send_form.style.display = "block"
                // 在消息区域显示连接成功的提示
                messages.innerHTML += "<p><em>You are connected.</em></p>"
            }

            // 当WebSocket接收到服务器端发送的消息时，执行以下操作
            ws.onmessage = function(event) {
                // 在消息区域显示服务器端发送的消息
                messages.innerHTML += "<p>" + event.data + "</p>"
            }

            // 当连接表单被提交时，执行以下操作
            form.onsubmit = function(event) {
                // 阻止表单的默认行为（刷新页面）
                event.preventDefault()
                // 获取用户输入的用户名
                var name = username.value
                // 如果用户名不为空，发送用户名给服务器端
                if (name) {
                    ws.send(name)
                }
            }

            // 当发送表单被提交时，执行以下操作
            send_form.onsubmit = function(event) {
                // 阻止表单的默认行为（刷新页面）
                event.preventDefault()
                // 获取用户输入的消息内容
                var text = message.value
                // 如果消息内容不为空，发送消息给服务器端
                if (text) {
                    ws.send(text)
                    // 清空消息输入框
                    message.value = ""
                }
            }
        </script>
    </body>
</html>
"""


# 定义一个路由，返回HTML文档作为响应
@app.get("/")
async def get():
    return HTMLResponse(html)


# 定义一个WebSocket路由，用于处理WebSocket连接和通信
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # 接受WebSocket连接请求
    await websocket.accept()
    # 从WebSocket中接收用户名，并将其作为键添加到连接字典中，值为WebSocket对象本身
    name = await websocket.receive_text()
    connections[name] = websocket
    # 向所有已连接的客户端发送欢迎消息
    for connection in connections.values():
        await connection.send_text(f"{name} joined the chat.")
    # 循环接收和发送消息，直到连接断开或发生异常
    try:
        while True:
            # 从WebSocket中接收消息内容
            text = await websocket.receive_text()
            # 向所有已连接的客户端发送消息内容，附上发送者的用户名
            for connection in connections.values():
                await connection.send_text(f"{name}: {text}")
    except:
        # 如果发生异常或连接断开，从连接字典中移除该WebSocket对象，并向其他客户端发送离开消息
        del connections[name]
        for connection in connections.values():
            await connection.send_text(f"{name} left the chat.")


@app.get("/a")
async def a():
    for connection in connections.values():
        await connection.send_text("这样子吗")
    return "11"
